package agkey

import (
	crand "crypto/rand"
	"encoding/asn1"
	"errors"
	"io"
	"log"
	"math/rand"

	ko "gitee.com/tgodfather/utility/crypt/options"
	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/sm3"
	"github.com/tjfoc/gmsm/sm4"
)

var (
	ERR_AGKEY_CAP_UNSUPPORT      = errors.New("unsupport capablity")
	ERR_AGKEY_NOT_FOUND          = errors.New("key not found")
	ERR_AGKEY_INVALID_CIPHERTEXT = errors.New("invalid cipher text")
	ERR_AGKEY_INVALID_PARAM      = errors.New("invalid params")
	ERR_AGKEY_UNEXCEPTED_KEY     = errors.New("unexpected key")
	ERR_AGKEY_KEY_UNSUPPORT      = errors.New("unsupport key")
)

// type ciphertext struct{
// 	Kid []byte
// 	CipherText []byte
// }

type asnitem struct {
	Kid    string
	KType  ko.KeyType
	Cipher string
}
type item struct {
	asnitem
	CKey any //`asn1:"optional,omitempty"`
	//`asn1:"optional,explicit"`
	//`asn1:"-"`
}

type agkey struct {
	keyMap    map[string](item)
	kidCapMap map[int]([]string) // [ktype]([]kid)
}

func newAgKey() *agkey {
	return &agkey{
		keyMap:    make(map[string]item),
		kidCapMap: make(map[int][]string),
	}
}

func New() *agkey {
	return newAgKey()
}

func (k *agkey) Load(data []byte, o ...ko.KeyOption) (err error) {
	// 如果有 io.reader, 以io.reader 为主， data忽略
	opts := ko.NewKeyOptions(o...)
	rest := data
	if opts.Reader != nil {
		rest, err = io.ReadAll(opts.Reader)
		if err != nil {
			return err
		}
	}

	// rest := body
	for {
		nit := &item{}

		rest, err = asn1.Unmarshal(rest, &nit.asnitem)
		// log.Println("========", len(rest), err, nit)

		if err != nil {
			break
		}
		switch nit.KType {
		case ko.KT_SM2_PRI:
			a, priErr := func_dec_pri(nit.Cipher)
			nit.CKey = a
			if priErr != nil {
				log.Println("load pri from hex", priErr, len(nit.Cipher))
			}
		case ko.KT_SM2_PUB:
			a, pubErr := func_dec_pub(nit.Cipher)
			nit.CKey = a
			if pubErr != nil {
				log.Println("load pub from hex", pubErr, len(nit.Cipher))
			}
		}
		k.Add(*nit)
		// log.Printf("newKey keylen[%d] enc[%d] dec[%d] sign[%d] verify[%d] \n",
		// 	k.Length(),
		// 	k.CapLength(ko.CT_ENCRYPT),
		// 	k.CapLength(ko.CT_DECRYPT),
		// 	k.CapLength(ko.CT_SIGN),
		// 	k.CapLength(ko.CT_VERIFY))
		if len(rest) <= 0 {
			break
		}
	}

	return err
}

type handler func(d []byte) error

func (k *agkey) Save(o ...ko.KeyOption) (result []byte, err error) {
	// 如果有 io.writer， 结果result就为nil
	opts := ko.NewKeyOptions(o...)
	var wh handler
	if opts.Writer != nil {
		result = nil
		wh = func(d []byte) error {
			_, weer := opts.Writer.Write(d)
			return weer
		}
	} else {
		wh = func(d []byte) error {
			result = append(result, d...)
			return nil
		}
	}

	for i, it := range k.keyMap {
		if d, mErr := asn1.Marshal(it.asnitem); mErr != nil {
			log.Println("agkey save marshal error", mErr, i)
			err = mErr
			break
		} else {
			if wErr := wh(d); wErr != nil {
				log.Println("agkey save error", wErr, i)
				err = wErr
				break
			}
		}
	}

	return result, err
}

func (k *agkey) Sign(data []byte, o ...ko.KeyOption) ([]byte, error) {

	if it := k.selectKey(ko.CT_SIGN); it != nil {
		if signatue, err := it.Sign(data, o...); err != nil {
			return nil, err
		} else {
			//增加keyid
			return asn1.Marshal(envelope{
				Kid:        []byte(it.Kid),
				Signature:  signatue,
				EncKey:     []byte(" "),
				CipherText: []byte(" "),
			})
		}
	}

	return nil, ERR_AGKEY_CAP_UNSUPPORT
}
func (k *agkey) Verify(data []byte, signature []byte, o ...ko.KeyOption) (bool, error) {
	env := &envelope{}
	if _, err := asn1.Unmarshal(signature, env); err != nil {
		return false, err
	} else {
		if it, bExist := k.keyMap[string(env.Kid)]; bExist {
			return it.Verify(data, env.Signature)
		} else {
			return false, ERR_AGKEY_NOT_FOUND
		}
	}

	// // if it := k.selectKey(ko.CT_VERIFY); it != nil {
	// // 	it.Verify(data, signature, o...)
	// // }
	// return false, ERR_AGKEY_CAP_UNSUPPORT
}

func (k *agkey) Encrypt(data []byte, o ...ko.KeyOption) ([]byte, error) {
	if it := k.selectKey(ko.CT_ENCRYPT); it != nil {

		// 生成随机sm4密钥, 长度为16
		buf := make([]byte, 32)
		crand.Read(buf)
		sm4Key := sm3.Sm3Sum(buf)[:16] // sm3Sum 结果长度为32个字节， sm4key 要求为16个字节

		// 用kid对应 的公钥加密 sm4密钥
		if sm4KeyCipher, err := it.Encrypt(sm4Key, o...); err != nil {
			return nil, err
		} else {
			//使用sm4key 加密 数据
			// log.Println("Encrypt sm4key ", sm4Key)
			if result, err := sm4.Sm4Ecb(sm4Key, data, true); err != nil {
				return nil, err
			} else {
				// log.Println("Encrypt  data", data, result)
				return asn1.Marshal(envelope{
					Kid:        []byte(it.Kid),
					Signature:  []byte(" "),
					EncKey:     sm4KeyCipher,
					CipherText: result,
				})
			}
		}
	}
	return nil, ERR_AGKEY_CAP_UNSUPPORT
}

func (k *agkey) Decrypt(data []byte, o ...ko.KeyOption) ([]byte, error) {
	env := &envelope{}
	if _, err := asn1.Unmarshal(data, env); err != nil {
		return nil, err
	} else {
		if it, bExist := k.keyMap[string(env.Kid)]; bExist {
			//解密 密钥
			if sm4Key, err := it.Decrypt(env.EncKey); err != nil {
				return nil, err
			} else {
				return sm4.Sm4Ecb(sm4Key, env.CipherText, false)
			}
		} else {
			return nil, ERR_AGKEY_NOT_FOUND
		}
	}

	// return nil, ERR_AGKEY_CAP_UNSUPPORT
}

//	func (k *agkey) itemDel(it item, o ...crypt.KeyOption) ([]byte, error) {
//		return nil, nil
//	}
//
// /////////////////////////////////
func (k *agkey) Add(it item, o ...ko.KeyOption) error {
	k.keyMap[it.Kid] = it
	//建立 相同CapType 的kid数组, 再建立数组的map， 方便定位各种CapType的key
	cap_of_item := func(kt ko.KeyType) []ko.CapType {
		switch kt {
		case ko.KT_SM2_PRI:
			return []ko.CapType{ko.CT_SIGN, ko.CT_DECRYPT, ko.CT_VERIFY, ko.CT_ENCRYPT}
		case ko.KT_SM2_PUB:
			return []ko.CapType{ko.CT_VERIFY, ko.CT_ENCRYPT}
		case ko.KT_SM4:
			return []ko.CapType{ko.CT_DECRYPT, ko.CT_ENCRYPT}
		}
		return []ko.CapType{}
	}

	caplist := cap_of_item(it.KType)
	for _, cap := range caplist {
		var kidA []string
		if ka, bExist := k.kidCapMap[int(cap)]; bExist {
			kidA = ka
		} else {
			kidA = make([]string, 0)
		}

		kidA = append(kidA, it.Kid)
		k.kidCapMap[int(cap)] = kidA
	}
	return nil
}

func (k *agkey) selectKey(cap ko.CapType) *item {
	// k.keyMap[it.Kid] = it
	if ka, bExist := k.kidCapMap[int(cap)]; bExist {

		knum := len(ka)
		if knum > 0 {
			kid := ka[rand.Intn(knum)]
			// log.Println("selectKey ", cap, len(ka), kid)
			if it, kexist := k.keyMap[kid]; kexist {
				return &it
			}
		}
	}

	return nil
}

func (k *agkey) Length() int {
	return len(k.keyMap)
}

func (k *agkey) CapLength(cap ko.CapType) int {
	return len(k.kidCapMap[int(cap)])
}

// item

func (it *item) Sign(data []byte, o ...ko.KeyOption) ([]byte, error) {
	if prikey := it.CKey.(*sm2.PrivateKey); prikey != nil {
		return prikey.Sign(crand.Reader, data, nil)
	}
	return nil, ERR_AGKEY_CAP_UNSUPPORT
}
func (it *item) Verify(data []byte, signature []byte, o ...ko.KeyOption) (bool, error) {
	var pubkey *sm2.PublicKey
	switch it.CKey.(type) {
	case (*sm2.PublicKey):
		pubkey = it.CKey.(*sm2.PublicKey)
	case (*sm2.PrivateKey):
		pubkey = &it.CKey.(*sm2.PrivateKey).PublicKey
	default:
		pubkey = nil
	}

	if pubkey != nil {
		isValid := pubkey.Verify(data, signature)
		return isValid, nil
	}

	return false, ERR_AGKEY_CAP_UNSUPPORT
}

func (it *item) Encrypt(data []byte, o ...ko.KeyOption) ([]byte, error) {
	var pubkey *sm2.PublicKey
	switch it.CKey.(type) {
	case (*sm2.PublicKey):
		pubkey = it.CKey.(*sm2.PublicKey)
	case (*sm2.PrivateKey):
		pubkey = &it.CKey.(*sm2.PrivateKey).PublicKey
	default:
		pubkey = nil
	}

	if pubkey != nil {
		return pubkey.EncryptAsn1(data, crand.Reader)
	}
	return nil, ERR_AGKEY_CAP_UNSUPPORT
}

func (it *item) Decrypt(data []byte, o ...ko.KeyOption) ([]byte, error) {
	if prikey := it.CKey.(*sm2.PrivateKey); prikey != nil {
		return prikey.DecryptAsn1(data)
	}
	return nil, ERR_AGKEY_CAP_UNSUPPORT
}
